#!/bin/sh

#
# Script which configures the grub.cfg to use an updated root index
#

set -o errexit

# shellcheck disable=SC1091
. /usr/sbin/balena-config-vars
# shellcheck disable=SC1091
. /usr/libexec/os-helpers-logging
# shellcheck disable=SC1091
. /usr/libexec/os-helpers-fs

verifyMd5sum() {
  if [ "$1" = "$3" ] ; then return 1; fi
}

getSize() {
  echo $(du -sb "$1" | awk '{print $1}')
}

DURING_UPDATE=${DURING_UPDATE:-0}

if [ "$DURING_UPDATE" = "1" ]; then
    SYSROOT="/mnt/sysroot/inactive"
else
    SYSROOT="/mnt/sysroot/active"
fi

new_part=$(findmnt --noheadings --canonicalize --output SOURCE $SYSROOT -t ext4)
enc_part=0
if is_part_encrypted "${new_part}"; then
    blockdev=$(lsblk -sJ "${new_part}" | jq -r '.blockdevices[].children[].name')
    enc_part=1
else
    blockdev=$(lsblk -sJ "${new_part}" | jq -r '.blockdevices[].name')
fi
new_part_idx=$(cat "/sys/class/block/$blockdev/partition")
new_part_label=$(blkid "$new_part" | awk '{print $2}' | cut -d'"' -f 2)

info "Switching root partition to $new_part_label..."

# flash legacy grub only if we do not support UEFI
if [ ! -d /sys/firmware/efi ] ; then
    # Remove EFI
    rm -rf "$BALENA_BOOT_MOUNTPOINT/EFI" || true

    MBR=512
    reservedMBR=66
    sourcePath="/resin-boot/grub/"
    device="/dev/$(findmnt --noheadings --canonicalize --output SOURCE /mnt/sysroot/active | xargs lsblk -no pkname)"

    firstBootloader=boot.img
    secondBootloader=core.img

    md5sum_diskFirstBootloader=$(dd if=$device bs=1 count=$(getSize "$sourcePath$firstBootloader") | md5sum)
    md5sum_diskSecondBootloader=$(dd if=$device skip=$MBR bs=1 count=$(getSize "$sourcePath$secondBootloader") | md5sum)

    if verifyMd5sum $(md5sum $sourcePath$firstBootloader) $md5sum_diskFirstBootloader; then
        dd if="$sourcePath$firstBootloader" of="$device" conv=fdatasync bs=1 count=$(expr $MBR - $reservedMBR)
    fi
    if verifyMd5sum $(md5sum $sourcePath$secondBootloader) $md5sum_diskSecondBootloader; then
        dd if="$sourcePath$secondBootloader" of="$device" conv=fdatasync bs=1 seek=$MBR
    fi
else
    # Remove legacy grub
    rm -rf "$BALENA_BOOT_MOUNTPOINT/grub" || true

    # Put the secure boot GRUB variant in place or remove it when applicable
    EFI_BOOT_DIR="${BALENA_BOOT_MOUNTPOINT}/EFI/BOOT/"
    for EFI_BINARY in "${EFI_BOOT_DIR}/"*.efi; do
        if [ "${EFI_BINARY}" = "${EFI_BOOT_DIR}/*.efi" ]; then
            # Nothing matched the glob
            break
        fi

        EFI_BINARY_SECUREBOOT="${EFI_BINARY}.secureboot"
        if [ "${enc_part}" = "1" ]; then
            # If this fails it is a fatal error - the secure boot variant
            # must be in place when LUKS is enabled
            mv "${EFI_BINARY_SECUREBOOT}" "${EFI_BINARY}"
        else
            # This can safely fail if the secure boot variant does not exist
            rm -f "${EFI_BINARY_SECUREBOOT}" || :
        fi
    done
fi

grub_cfg=$(find -L $BALENA_BOOT_MOUNTPOINT -name grub.cfg)
grub_luks_cfg=$(find -L $BALENA_BOOT_MOUNTPOINT -name grub-luks.cfg)
grub_env=$(find -L $BALENA_BOOT_MOUNTPOINT -name grubenv)

if [ -n "$grub_luks_cfg" ]; then
    if [ "${enc_part}" = "1" ]; then
        "${MV}" "$grub_luks_cfg" "$grub_cfg"
        "${MV}" "$grub_luks_cfg.sig" "$grub_cfg.sig"
    else
        rm -f "$grub_luks_cfg" "$grub_luks_cfg.sig"
    fi
fi

tmpfile=$(mktemp)
if  "${CAT}" "$grub_env" | grep -q upgrade_available; then
    info "Automated Rollback support in grub.cfg detected"
    "${CAT}" "$grub_env" | sed -e "s#resin_root_part=.*#resin_root_part="$new_part_idx"#g" -e "s#upgrade_available=.*#upgrade_available="$DURING_UPDATE"#g" -e "s#bootcount=.*#bootcount=0#g" | "${WR}" "${grub_env}"
else
    info "Automated Rollback is not supported by grub config file for this device"
    "${CAT}" "$grub_cfg" | sed "s/resin-root./${new_part_label}/" | "${WR}" "${grub_cfg}"
fi
info "Switch root done."
